#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <cuda_runtime.h>
#include <fstream>
#include <iostream>
#include <memory>

// 自定义日志类（必须实现）
class Logger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        // 只打印警告及以上信息
        if (severity <= Severity::kWARNING)
            std::cout << "[TensorRT] " << msg << std::endl;
    }
} gLogger;

// 将 ONNX 转换为 TensorRT 引擎并保存
bool onnxToTensorRT(const std::string& onnxPath, const std::string& enginePath, 
                    bool useFP16 = false, int maxBatchSize = 1) {
    // 1. 创建 Builder
    auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(gLogger));
    if (!builder) {
        std::cerr << "创建 builder 失败！" << std::endl;
        return false;
    }

    // 2. 创建网络（显式批处理）
    auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(
        builder->createNetworkV2(1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH))
    );

    // 3. 创建 ONNX Parser 并解析模型
    auto parser = std::unique_ptr<nvonnxparser::IParser>(
        nvonnxparser::createParser(*network, gLogger)
    );
    if (!parser->parseFromFile(onnxPath.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kWARNING))) {
        std::cerr << "解析 ONNX 文件失败！" << std::endl;
        return false;
    }

    // 4. 配置 Builder 参数
    // builder->setMaxBatchSize(maxBatchSize);  // 已废弃，无需设置
    // config->setMaxWorkspaceSize(1ULL << 30); // 工作空间大小（1GB）
    if (useFP16) {
        config->setFlag(nvinfer1::BuilderFlag::kFP16);  // 直接设置FP16
    }

    // 5. 构建 TensorRT 引擎（先序列化，再反序列化）
    auto serializedEngine = std::unique_ptr<nvinfer1::IHostMemory>(
        builder->buildSerializedNetwork(*network, *config)
    );
    if (!serializedEngine) {
        std::cerr << "序列化网络失败！" << std::endl;
        return false;
    }

    // 反序列化为 ICudaEngine
    auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(gLogger));
    auto engine = std::unique_ptr<nvinfer1::ICudaEngine>(
        runtime->deserializeCudaEngine(serializedEngine->data(), serializedEngine->size())
    );
    if (!engine) {
        std::cerr << "构建 TensorRT 引擎失败！" << std::endl;
        return false;
    }

    // 6. 保存序列化引擎到文件
    std::ofstream file(enginePath, std::ios::binary);
    file.write(reinterpret_cast<const char*>(serializedEngine->data()), serializedEngine->size());
    file.close();

    std::cout << "序列化并保存 TensorRT 引擎成功: " << enginePath << std::endl;
    return true;
}

int main() {
    std::string onnxPath = "/home/nano/Documents/rm_f_03/yolo_tensorRT/test/src/yolov8s.onnx";      // 输入 ONNX 文件路径
    std::string enginePath = "yolov8s.engine";  // 输出 TensorRT 引擎路径
    bool useFP16 = true;                        // 是否启用 FP16
    int maxBatchSize = 1;                       // 最大批大小

    if (onnxToTensorRT(onnxPath, enginePath, useFP16, maxBatchSize)) {
        std::cout << "ONNX 转 TensorRT 引擎成功！" << std::endl;
    } else {
        std::cerr << "ONNX 转 TensorRT 引擎失败！" << std::endl;
        return -1;
    }
    return 0;
}